// Package checksum checksum/manager.go
package checksum

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"px-table-checksum-go/compare"
	"px-table-checksum-go/config"
	"px-table-checksum-go/db"
	"px-table-checksum-go/models"
	"px-table-checksum-go/queue"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	log "github.com/sirupsen/logrus"
)

// 定义错误常量，便于统一错误处理
var (
	// ErrInvalidDBConnection 表示数据库连接无效
	ErrInvalidDBConnection = errors.New("无效的数据库连接")

	// ErrTableProcessingFailed 表示表处理失败
	ErrTableProcessingFailed = errors.New("表处理失败")

	// 默认输出目录
	outputDir = "./fixsql"
)

// RunChecksumAndCompare 运行校验和比较
func RunChecksumAndCompare() error {
	// 在主函数中设置超时30分钟的上下文
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
	defer cancel()

	return RunChecksumAndCompareWithContext(ctx)
}

// RunChecksumAndCompareWithContext 支持上下文的校验和比较函数
func RunChecksumAndCompareWithContext(ctx context.Context) error {
	startTime := time.Now()
	log.Debug("进入RunChecksumAndCompareWithContext函数")

	// 获取配置并检查有效性
	cfg := config.GlobalConfig
	options := cfg.Options
	// 检查Options是否为零值
	log.Debugf("配置选项: Mode=%d, ChunkSize=%d, RunFixSQL=%t, GenerateFixSQL=%t, SourceDBID=%s, TargetDBID=%s",
		options.Mode, options.ChunkSize, options.RunFixSQL, options.GenerateFixSQL,
		options.SourceDBID, options.TargetDBID)

	if options.ChunkSize <= 0 {
		log.Error("配置选项无效：ChunkSize必须大于0")
		return errors.New("配置选项无效：ChunkSize必须大于0")
	}

	dbManager := db.GetDBManager()
	if dbManager == nil {
		log.Error("无法获取数据库管理器")
		return errors.New("无法获取数据库管理器")
	}

	checksumDB := dbManager.GetChecksumDB()
	sourceDB := dbManager.GetSourceDB()

	// 使用sourceDB变量避免未使用警告
	if sourceDB == nil || checksumDB == nil {
		log.Errorf("数据库连接无效: sourceDB=%v, checksumDB=%v", sourceDB != nil, checksumDB != nil)
		return fmt.Errorf("%w: 数据库连接无效", ErrInvalidDBConnection)
	}

	// 执行Checksum计算
	if options.Mode == 0 || options.Mode == 1 {
		log.Info("模式 0 或 1: 开始执行 Checksum 计算...")

		// 准备checksum表
		if err := db.PrepareChecksumTable(checksumDB); err != nil {
			log.Errorf("准备结果表失败: %v", err)
			return fmt.Errorf("准备结果表失败: %w", err)
		}

		// 实现校验计算逻辑
		log.Debug("开始调用runChecksumCalculationWithContext执行校验计算...")
		if err := runChecksumCalculationWithContext(ctx, dbManager); err != nil {
			log.Errorf("校验计算失败: %v", err)
			return fmt.Errorf("校验计算失败: %w", err)
		}
		log.Debug("校验计算完成")
	} else {
		log.Debugf("跳过校验计算阶段(当前模式: %d)", options.Mode)
	}

	// 执行校验结果比较和修复
	if options.Mode == 1 || options.Mode == 2 {
		log.Info("模式 1 或 2: 开始比较和修复...")

		// 添加详细日志，便于调试
		log.Debugf("开始调用runCompareAndRepairWithContext执行比较和修复...")
		if err := runCompareAndRepairWithContext(ctx, dbManager); err != nil {
			log.Errorf("比较和修复过程失败: %v", err)
			return fmt.Errorf("比较和修复过程失败: %w", err)
		}
		log.Debug("比较和修复完成")
	} else {
		log.Debugf("跳过比较和修复阶段(当前模式: %d)", options.Mode)
	}

	log.Infof("全部操作完成，总耗时: %s", time.Since(startTime))
	log.Info("比较和修复过程已全部完成，应用即将退出")
	return nil
}

// runCompareAndRepairWithContext 执行表比较和修复
func runCompareAndRepairWithContext(ctx context.Context, dbManager *db.DBConnectionManager) error {
	log.Debugf("开始执行 runCompareAndRepairWithContext 函数")
	cfg := config.GlobalConfig
	checksumDB := dbManager.GetChecksumDB()
	sourceDB := dbManager.GetSourceDB()

	if checksumDB == nil {
		log.Error("Checksum DB连接为空")
		return fmt.Errorf("%w: Checksum DB 连接无效", ErrInvalidDBConnection)
	}

	log.Debugf("配置选项: Mode=%d, GenerateFixSQL=%t, RunFixSQL=%t",
		cfg.Options.Mode, cfg.Options.GenerateFixSQL, cfg.Options.RunFixSQL)

	// 记录比较结果
	allConsistent := true
	totalDiffChunksFound := 0
	totalFixSQLsGenerated := 0

	// 获取配置的表总数
	tablesTotal := countConfiguredTables(cfg.TablesCheck)
	log.Infof("配置中共有 %d 个表", tablesTotal)

	if tablesTotal == 0 {
		log.Warn("配置中没有配置任何表，比较过程将跳过")
		return nil
	}

	// 构建表列表
	tables := buildTableList(cfg.TablesCheck)
	log.Debugf("已构建表列表，共 %d 个表", len(tables))

	// 处理每个表
	log.Infof("开始逐个处理 %d 个表的比较...", len(tables))
	for tableIndex, tableInfo := range tables {
		// 检查上下文是否已取消
		select {
		case <-ctx.Done():
			log.Debugf("比较操作被取消: %v", ctx.Err())
			return fmt.Errorf("比较操作被取消: %w", ctx.Err())
		default:
			// 继续处理
		}

		schemaName := tableInfo.Schema
		tableName := tableInfo.Table

		log.Infof("\n==================== 比较表 (%d/%d): %s.%s ====================",
			tableIndex+1, len(tables), schemaName, tableName)

		// 为每个表创建超时子上下文
		tableCtx, tableCancel := context.WithTimeout(ctx, 5*time.Minute)
		// 确保在当前表处理完成时取消上下文，而不是在整个函数结束时
		// 这将防止上下文堆积，导致程序卡住
		// defer tableCancel() - 移到下面的作用域中，不要放在这里

		// 使用一个立即执行的函数来限制tableCancel的作用域
		func() {
			// 确保在此函数结束时取消当前表的上下文
			defer tableCancel()

			// 检查校验数据是否存在
			log.Debugf("检查表 %s.%s 的校验数据是否存在", schemaName, tableName)
			sourceExists, targetExists, err := checkTableDataExists(checksumDB, schemaName, tableName, cfg.Options.SourceDBID, cfg.Options.TargetDBID)
			if err != nil {
				log.Debugf("检查表 %s.%s 的校验数据时出错: %v", schemaName, tableName, err)
			}

			log.Debugf("表 %s.%s 的校验数据存在情况: 源端=%t, 目标端=%t", schemaName, tableName, sourceExists, targetExists)

			if !sourceExists || !targetExists {
				log.Warnf("跳过表 %s.%s 的比较，因为校验数据不完整(源:%t,目标:%t)", schemaName, tableName, sourceExists, targetExists)
				return // 取消子上下文通过defer自动执行
			}

			// 使用比较服务，设置超时
			log.Infof("开始获取表 %s.%s 的差异块...", schemaName, tableName)
			log.Debugf("调用 compare.NewCompareChecksums 比较校验和...")
			diffChunks, err := compare.NewCompareChecksums(schemaName, tableName)
			if err != nil {
				log.Errorf("错误: 比较 Checksum 失败 (%s.%s): %v", schemaName, tableName, err)
				allConsistent = false
				return // 取消子上下文通过defer自动执行
			}

			uniqueDiffChunkCount := len(diffChunks)
			log.Infof("找到 %d 个不一致的数据块", uniqueDiffChunkCount)
			log.Debugf("diffChunks数组长度: %d", uniqueDiffChunkCount)

			if uniqueDiffChunkCount == 0 {
				log.Infof("表 %s.%s 数据一致", schemaName, tableName)
				return // 取消子上下文通过defer自动执行
			} else {
				log.Infof("表 %s.%s 发现 %d 个不一致的数据块", schemaName, tableName, uniqueDiffChunkCount)
				allConsistent = false
				totalDiffChunksFound += uniqueDiffChunkCount

				// 生成修复SQL的代码
				if (cfg.Options.GenerateFixSQL || cfg.Options.RunFixSQL) && sourceDB != nil {
					log.Infof("开始为表 %s.%s 生成修复SQL（共%d个数据块）...", schemaName, tableName, uniqueDiffChunkCount)
					log.Debugf("GenerateFixSQL=%t, RunFixSQL=%t, sourceDB不为空",
						cfg.Options.GenerateFixSQL, cfg.Options.RunFixSQL)

					// 设置生成SQL的超时上下文，给予更充足的时间
					sqlGenTimeout := 10 * time.Minute
					log.Debugf("设置SQL生成超时为%v", sqlGenTimeout)
					sqlGenCtx, sqlGenCancel := context.WithTimeout(tableCtx, sqlGenTimeout)
					// 确保SQL生成上下文在此函数结束时被取消
					defer sqlGenCancel()

					log.Debugf("调用 generateFixSQLsForTable 生成修复SQL...")
					sqlStartTime := time.Now()
					sqlCount, err := generateFixSQLsForTable(sqlGenCtx, sourceDB, checksumDB, schemaName, tableName, diffChunks)
					sqlElapsedTime := time.Since(sqlStartTime)

					if err != nil {
						// 检查是否是因为超时导致的错误
						if errors.Is(err, context.DeadlineExceeded) {
							log.Errorf("为表 %s.%s 生成修复SQL超时，请检查表结构或数据量 (已耗时: %v)",
								schemaName, tableName, sqlElapsedTime)
						} else if errors.Is(err, context.Canceled) {
							log.Warnf("为表 %s.%s 生成修复SQL被取消 (已耗时: %v)",
								schemaName, tableName, sqlElapsedTime)
						} else {
							log.Errorf("为表 %s.%s 生成修复SQL失败: %v (已耗时: %v)",
								schemaName, tableName, err, sqlElapsedTime)
						}
					} else {
						totalFixSQLsGenerated += sqlCount
						log.Infof("表 %s.%s 成功生成 %d 条修复SQL，总耗时: %v",
							schemaName, tableName, sqlCount, sqlElapsedTime)
					}

					log.Debugf("SQL生成完成")
				} else {
					if sourceDB == nil {
						log.Debugf("sourceDB为空，跳过生成修复SQL")
					}
					if !cfg.Options.GenerateFixSQL && !cfg.Options.RunFixSQL {
						log.Debugf("未启用生成或执行修复SQL选项，跳过")
					}
				}
			}

			log.Infof("==================== 表 %s.%s 比较完成 ====================\n", schemaName, tableName)
			log.Debugf("取消表 %s.%s 的处理上下文", schemaName, tableName)
			// tableCancel 会通过 defer 自动执行
		}()

		log.Debugf("表 %s.%s 处理器已结束, 继续处理下一个表", schemaName, tableName)
	}

	log.Info("所有表比较完成, 开始输出最终结果...")
	// 输出最终结果
	log.Info("-------------------- 最终比较结果 --------------------")
	if allConsistent {
		log.Info("所有检查的表数据一致")
	} else {
		log.Infof("检查完成，发现数据不一致。总计 %d 个不一致的数据块", totalDiffChunksFound)
		if totalFixSQLsGenerated > 0 {
			log.Infof("总共生成了 %d 条修复 SQL 语句", totalFixSQLsGenerated)
		} else if cfg.Options.GenerateFixSQL {
			log.Info("未生成任何修复 SQL 语句")
		}
	}
	log.Info("------------------------------------------------------")
	log.Info("比较和修复阶段已完全结束") // 添加明确的结束标记
	log.Debug("runCompareAndRepairWithContext函数执行完毕, 准备退出")

	return nil
}

// generateFixSQLsForTable 为表生成修复SQL
func generateFixSQLsForTable(ctx context.Context, sourceDB *sql.DB, checksumDB *sql.DB, schemaName, tableName string, diffChunks []*models.DiffChunk) (int, error) {
	cfg := config.GlobalConfig
	logger := log.WithFields(log.Fields{
		"schemaName": schemaName,
		"tableName":  tableName,
		"func":       "generateFixSQLsForTable",
	})
	logger.Info("准备生成修复SQL")
	logger.Debugf("发现 %d 个差异块需要处理", len(diffChunks))

	// 检查上下文是否已取消
	select {
	case <-ctx.Done():
		logger.Debugf("生成SQL上下文已取消: %v", ctx.Err())
		return 0, ctx.Err()
	default:
		// 继续处理
	}

	// 获取表的键信息
	logger.Debugf("开始获取表 %s.%s 的键信息", schemaName, tableName)
	keyInfo, err := db.GetBestUniqueKey(sourceDB, schemaName, tableName)
	if err != nil {
		logger.Errorf("获取表 %s.%s 的键信息失败: %v", schemaName, tableName, err)
		return 0, fmt.Errorf("获取表 %s.%s 的键信息失败: %w", schemaName, tableName, err)
	}
	logger.Debugf("获取到键信息: 键列=%v, 是否主键=%v", keyInfo.Columns, keyInfo.IsPrimary)

	// 检查键信息的有效性
	if keyInfo == nil || len(keyInfo.Columns) == 0 {
		logger.Error("获取到的键信息无效: 键列为空或为nil")
		return 0, fmt.Errorf("表 %s.%s 没有有效的主键或唯一键", schemaName, tableName)
	}

	// 获取表的所有列信息
	logger.Debugf("开始获取表 %s.%s 的列信息", schemaName, tableName)
	allColsInfo, err := db.GetTableColumnsAndTypes(sourceDB, schemaName, tableName)
	if err != nil {
		logger.Warnf("获取表 %s.%s 的列信息失败: %v, 将使用空的列信息列表继续", schemaName, tableName, err)
		allColsInfo = []db.ColumnInfo{} // 使用空列表，稍后可能会降级处理
	} else {
		logger.Debugf("获取到表的 %d 个列信息", len(allColsInfo))
		// 打印前几个列信息
		for i, col := range allColsInfo {
			if i < 3 || log.GetLevel() >= log.DebugLevel {
				logger.Debugf("列信息 #%d: 名称=%s, 类型=%s", i+1, col.Name, col.DataType)
			}
		}
	}

	// 清理旧的修复SQL文件
	if cfg.Options.GenerateFixSQL {
		logger.Debugf("开始清理表 %s.%s 的旧修复SQL文件", schemaName, tableName)
		cleanOldFixSQLFiles(schemaName, tableName)
	}

	// 定义输出目录
	outputDir := "./fixsql"
	logger.Debugf("修复SQL输出目录: %s", outputDir)

	// 为每个差异块生成修复SQL
	totalSQLCount := 0

	// 无需信号量，简化为串行处理
	for _, chunk := range diffChunks {
		// 简化上下文检查
		if ctx.Err() != nil {
			return totalSQLCount, ctx.Err()
		}

		// 找出具体行级别的差异
		logger.Debugf("开始调用 NewFindDifferencesInChunk 查找块 %d 的具体差异...", chunk.ChunkNo)
		startTime := time.Now()
		diffResult, err := compare.NewFindDifferencesInChunkWithContext(ctx, schemaName, tableName, chunk, keyInfo)
		elapsedTime := time.Since(startTime)
		logger.Debugf("查找块 %d 差异耗时: %v", chunk.ChunkNo, elapsedTime)

		if err != nil {
			// 处理错误情况
			logger.Errorf("查找块 %d 的具体差异失败: %v", chunk.ChunkNo, err)
			continue
		}

		// 判断diffResult是否为空
		if diffResult == nil {
			logger.Errorf("块 %d 的差异结果为空, 跳过生成SQL", chunk.ChunkNo)
			continue
		}

		// 判断是否有实际差异
		hasDiff := diffResult.OnlyInSource > 0 || diffResult.OnlyInTarget > 0 || diffResult.DifferentRows > 0
		logger.Debugf("块 %d 差异详情: 源端独有=%d行, 目标端独有=%d行, 数据不同=%d行",
			chunk.ChunkNo, diffResult.OnlyInSource, diffResult.OnlyInTarget, diffResult.DifferentRows)

		if !hasDiff {
			logger.Debugf("块 %d 没有发现具体行差异，跳过生成SQL", chunk.ChunkNo)
			continue
		}

		// 详细记录行级差异数据
		logger.Debugf("块 %d 源端独有行数据数组长度: %d", chunk.ChunkNo, len(diffResult.SourceOnlyRows))
		if len(diffResult.SourceOnlyRows) > 0 && log.GetLevel() >= log.DebugLevel {
			sampleRow := diffResult.SourceOnlyRows[0]
			logger.Debugf("块 %d 源端独有行样例: %+v", chunk.ChunkNo, sampleRow)
		}

		logger.Debugf("块 %d 目标端独有行数据数组长度: %d", chunk.ChunkNo, len(diffResult.TargetOnlyRows))
		if len(diffResult.TargetOnlyRows) > 0 && log.GetLevel() >= log.DebugLevel {
			sampleRow := diffResult.TargetOnlyRows[0]
			logger.Debugf("块 %d 目标端独有行样例: %+v", chunk.ChunkNo, sampleRow)
		}

		logger.Debugf("块 %d 数据不同行差异数组长度: %d", chunk.ChunkNo, len(diffResult.DifferentDetails))
		if len(diffResult.DifferentDetails) > 0 && log.GetLevel() >= log.DebugLevel {
			sampleDiff := diffResult.DifferentDetails[0]
			logger.Debugf("块 %d 数据不同行样例: 键=%+v, 不同列数=%d",
				chunk.ChunkNo, sampleDiff.Key, len(sampleDiff.DiffColumns))
		}

		// 生成修复SQL
		logger.Debugf("开始为块 %d 生成修复SQL...", chunk.ChunkNo)
		sqlStartTime := time.Now()
		fixSQLs, err := compare.GenerateFixSQLs(schemaName, tableName, keyInfo, allColsInfo, diffResult)
		sqlElapsedTime := time.Since(sqlStartTime)
		logger.Debugf("生成块 %d 修复SQL耗时: %v", chunk.ChunkNo, sqlElapsedTime)

		if err != nil {
			logger.Errorf("生成修复SQL失败: %v", err)
			continue
		}

		// 统计SQL数量
		sqlCount := len(fixSQLs)
		logger.Debugf("块 %d 生成修复SQL完成，共 %d 条SQL语句", chunk.ChunkNo, sqlCount)
		if sqlCount == 0 {
			// 增加详细的错误诊断信息
			logger.Warnf("块 %d 有差异数据但未生成任何SQL，源端独有=%d(%d), 目标端独有=%d(%d), 数据不同=%d(%d)",
				chunk.ChunkNo,
				diffResult.OnlyInSource, len(diffResult.SourceOnlyRows),
				diffResult.OnlyInTarget, len(diffResult.TargetOnlyRows),
				diffResult.DifferentRows, len(diffResult.DifferentDetails))

			continue
		}

		totalSQLCount += sqlCount

		// 如果需要生成SQL文件
		if cfg.Options.GenerateFixSQL && sqlCount > 0 {
			// 检查上下文是否取消
			select {
			case <-ctx.Done():
				logger.Debugf("写入SQL文件前上下文已取消")
				return totalSQLCount, ctx.Err()
			default:
				// 继续处理
			}

			// 确保输出目录存在
			logger.Debugf("准备将SQL写入文件，确保输出目录存在: %s", outputDir)
			if err := os.MkdirAll(outputDir, 0755); err != nil {
				logger.Errorf("创建输出目录失败: %v", err)
				return totalSQLCount, ctx.Err()
			}

			// 构建输出文件路径
			timestamp := time.Now().Format("20060102-150405")
			sqlFilePath := filepath.Join(outputDir,
				fmt.Sprintf("datafix-%s.%s-chunk%d-%s.sql",
					schemaName, tableName, chunk.ChunkNo, timestamp))

			// 写入SQL到文件
			logger.Debugf("开始写入 %d 条SQL到文件: %s", sqlCount, sqlFilePath)
			writeStartTime := time.Now()
			if err := writeFixSQLToFile(sqlFilePath, fixSQLs); err != nil {
				logger.Errorf("写入SQL文件失败: %v", err)
			} else {
				writeElapsedTime := time.Since(writeStartTime)
				logger.Infof("为块 %d 生成了 %d 条SQL，已写入文件: %s (耗时: %v)",
					chunk.ChunkNo, sqlCount, sqlFilePath, writeElapsedTime)

				// 记录前几条SQL示例
				if log.GetLevel() >= log.DebugLevel && len(fixSQLs) > 0 {
					maxSamples := 3
					if len(fixSQLs) < maxSamples {
						maxSamples = len(fixSQLs)
					}
					for i := 0; i < maxSamples; i++ {
						logger.Debugf("SQL示例 #%d: %s", i+1, fixSQLs[i])
					}
				}
			}
		} else if sqlCount > 0 {
			logger.Infof("为块 %d 生成了 %d 条SQL", chunk.ChunkNo, sqlCount)
			if !cfg.Options.GenerateFixSQL {
				logger.Debugf("未启用GenerateFixSQL选项，不写入SQL文件")
			}
		}

		// 如果需要执行修复SQL
		if cfg.Options.RunFixSQL && sqlCount > 0 {
			logger.Warnf("配置指定需要执行修复SQL，但当前版本暂不支持自动执行")
		}

		blockElapsedTime := time.Since(startTime)
		logger.Infof("块 %d 处理完成，总耗时: %v", chunk.ChunkNo, blockElapsedTime)
	}

	return totalSQLCount, nil
}

// buildTableList 从配置中构建表列表
func buildTableList(tablesCheck map[string][]string) []tableInfo {
	tables := make([]tableInfo, 0)
	for schema, tableNames := range tablesCheck {
		for _, tableName := range tableNames {
			tables = append(tables, tableInfo{
				Schema: schema,
				Table:  tableName,
			})
		}
	}
	return tables
}

// runChecksumCalculationWithContext 支持上下文的校验计算
func runChecksumCalculationWithContext(ctx context.Context, dbManager *db.DBConnectionManager) error {
	// 获取数据库连接
	sourceDB := dbManager.GetSourceDB()
	targetDB := dbManager.GetTargetDB()
	checksumDB := dbManager.GetChecksumDB()

	if sourceDB == nil || targetDB == nil || checksumDB == nil {
		return fmt.Errorf("%w: 获取数据库连接失败", ErrInvalidDBConnection)
	}

	// 设置工作线程数量
	workerCount := 4 // 默认工作线程数量

	// 使用当前CPU核心数，但最多8个线程
	cpuCores := runtime.NumCPU()
	if cpuCores > 0 {
		workerCount = cpuCores
		if workerCount > 8 {
			// 限制最大工作线程数，避免过多并发导致数据库压力过大
			workerCount = 8
		}
	}
	log.Infof("校验计算启动中，目标端工作线程数: %d", workerCount) // 修改日志说明

	// 创建任务队列 (根据需要选择 BoundedQueue 或 RedisQueue)
	// 注意: 这里假设使用 BoundedQueue，如果使用 RedisQueue，请确保正确初始化
	taskQueue := createTaskQueue() // 假设 createTaskQueue 返回 *queue.BoundedQueue

	// 创建已完成表计数器 (由 Source Worker 更新)
	finishedTableCounter := &atomic.Int32{}

	// 统计要处理的表的总数
	var tablesCount int // 重命名以避免与 finishedTableCounter 混淆

	// 启动目标端工作线程
	var targetWg sync.WaitGroup
	targetWg.Add(workerCount)

	// 创建取消函数，在发生错误时可以通知所有goroutine退出
	// 注意: runChecksumCalculationWithContext 自身也可能接收一个 ctx，需要正确传递
	ctx, cancel := context.WithCancel(ctx) // 使用传入的 ctx 创建子上下文
	defer cancel()                         // 确保在函数退出时取消所有 goroutine

	// 启动工作者线程
	for i := 0; i < workerCount; i++ {
		go func(workerID int) {
			defer targetWg.Done() // 确保每个 goroutine 退出时调用 Done
			log.Debugf("启动目标端工作线程 #%d", workerID)
			// ***修改点: 不再传递 totalTables 给 Target Worker***
			RunTargetChecksumWorker(workerID, taskQueue)
		}(i)
	}

	// 启动源端处理
	startTime := time.Now()

	// 首先，获取要处理的表列表并启动 Source Workers
	// ***修改点: 直接使用 *sql.DB 传递，如果 Connection 只是简单包装***
	var errProcessTables error
	tablesCount, errProcessTables = processConfiguredTablesWithContext(ctx, sourceDB, checksumDB, taskQueue, finishedTableCounter)
	if errProcessTables != nil {
		log.Errorf("处理配置表错误: %v", errProcessTables)
		cancel()
		taskQueue.Close()
		targetWg.Wait()
		return fmt.Errorf("%w: %v", ErrTableProcessingFailed, errProcessTables)
	}

	if tablesCount == 0 {
		log.Warn("没有配置或找到需要处理的表")
		// 即使没有表，也需要正确关闭队列并等待 target workers 退出
		taskQueue.Close()
		targetWg.Wait()
		return nil // 没有表处理也是一种正常完成
	}

	log.Infof("已调度 %d 张表进行校验", tablesCount)

	// 等待所有 *Source* 表处理完成 (计数器只由 Source 更新)
	log.Debugf("等待 %d 张表由源端处理完成...", tablesCount)
	waitForCompletionWithContext(ctx, finishedTableCounter, tablesCount) // 等待计数达到 tablesCount
	log.Debug("所有源端工作线程已完成表的处理")

	// 关闭任务队列，通知目标端工作线程处理完剩余任务后退出
	// 这个操作必须在 waitForCompletionWithContext 之后，targetWg.Wait 之前
	log.Debug("所有表处理完成，关闭任务队列，通知目标端处理完剩余任务")
	taskQueue.Close()

	// 等待目标端工作线程完成处理队列中所有剩余的任务并退出
	log.Debug("等待所有目标端工作线程完成...")
	targetWg.Wait()
	log.Debug("所有目标端工作线程已退出")

	// 计算总耗时
	totalTime := time.Since(startTime)
	log.Infof("校验计算完成，总耗时: %v", totalTime)

	return nil
}

// processConfiguredTablesWithContext 支持上下文的表处理
// ***修改点: 参数类型改为 *sql.DB***
func processConfiguredTablesWithContext(ctx context.Context, sourceDB *sql.DB, checksumDB *sql.DB, taskQueue queue.Queue, finishedTableCounter *atomic.Int32) (int, error) {
	cfg := config.GlobalConfig
	tablesTotal := 0

	// 使用工作池处理多个表
	var workerWg sync.WaitGroup
	maxConcurrentTables := runtime.NumCPU() // 使用CPU核心数作为并发处理表的数量
	if maxConcurrentTables > 8 {
		maxConcurrentTables = 8 // 最多限制8个并发
	}
	if maxConcurrentTables < 1 {
		maxConcurrentTables = 1 // 至少有1个工作线程
	}
	log.Infof("源端表处理并发数: %d", maxConcurrentTables) // 添加日志

	// 创建表处理通道
	tableChan := make(chan tableInfo, maxConcurrentTables*2) // 增加缓冲区

	// 创建错误通道用于收集工作线程错误
	errChan := make(chan error, maxConcurrentTables)

	// 创建取消函数，用于在遇到错误时通知所有goroutine退出
	// 使用传入的 ctx 创建子上下文
	procCtx, cancel := context.WithCancel(ctx)
	defer cancel() // 确保此函数退出时取消内部 goroutine

	// 启动表处理工作线程 (Source Workers)
	for i := 0; i < maxConcurrentTables; i++ {
		workerWg.Add(1)
		go func(workerID int) {
			defer workerWg.Done()
			for {
				select {
				case <-procCtx.Done(): // 监听取消信号
					log.Debugf("源端工作线程 #%d 收到取消信号，退出", workerID)
					return
				case tableInfo, ok := <-tableChan:
					if !ok {
						log.Debugf("源端工作线程 #%d 表通道关闭，退出", workerID)
						return // 通道关闭，退出
					}
					// 处理单个表
					log.Infof("源端工作线程 #%d 开始处理表: %s.%s", workerID, tableInfo.Schema, tableInfo.Table)
					// ***修改点: 传递 *sql.DB 而不是 *db.Connection***
					// ***注意: RunSourceChecksumWorker 内部如果需要 Connection 包装器的方法，需要调整 RunSourceChecksumWorker 或此处的调用***
					// 这里假设 RunSourceChecksumWorker 可以直接使用 *sql.DB
					err := RunSourceChecksumWorker(sourceDB, checksumDB, taskQueue, finishedTableCounter, tableInfo.Schema, tableInfo.Table, procCtx) // 传递 procCtx
					if err != nil {
						// 检查是否是上下文取消错误
						if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
							log.Warnf("源端工作线程 #%d 处理表 %s.%s 被取消或超时: %v", workerID, tableInfo.Schema, tableInfo.Table, err)
						} else {
							log.Errorf("源端工作线程 #%d 处理表 %s.%s 失败: %v", workerID, tableInfo.Schema, tableInfo.Table, err)
							// 将错误发送到错误通道，并取消其他 worker
							select {
							case errChan <- fmt.Errorf("处理表 %s.%s 失败: %w", tableInfo.Schema, tableInfo.Table, err):
							case <-procCtx.Done(): // 如果已经被取消了，就不再发送错误
							}
							cancel() // 取消所有其他操作
						}
						// 即使出错或取消，也要继续处理下一个表或退出
						continue // 继续循环等待下一个任务或退出信号
					}
					log.Infof("源端工作线程 #%d 完成处理表: %s.%s", workerID, tableInfo.Schema, tableInfo.Table)
				}
			}
		}(i)
	}

	// 启动错误收集goroutine
	errCollected := make(chan struct{})
	var processingErr error
	go func() {
		defer close(errCollected)
		// 只接收第一个错误
		err, ok := <-errChan
		if ok {
			processingErr = err
			cancel() // 收到错误，立即取消所有 worker
		}
		// 读取完或通道关闭后退出
	}()

	// 遍历配置的表并加入处理队列
	tables := buildTableList(cfg.TablesCheck)
	log.Debugf("准备将 %d 个表加入处理通道...", len(tables))
	for _, table := range tables {
		select {
		case <-procCtx.Done(): // 如果在添加过程中被取消，则停止添加
			log.Warn("表处理任务添加过程中收到取消信号，停止添加剩余表")
			goto WaitForWorkers // 跳转到等待 worker 结束
		case tableChan <- table:
			tablesTotal++
		}
	}
	log.Debugf("已将 %d 个表加入处理通道", tablesTotal)

WaitForWorkers:
	// 关闭表处理通道，通知 worker 没有更多表了
	log.Debug("关闭表处理通道")
	close(tableChan)

	// 等待所有表工作线程完成
	log.Debug("等待所有源端工作线程完成...")
	workerWg.Wait()
	log.Debug("所有源端工作线程已完成")

	// 关闭错误通道（确保收集 goroutine 能退出）
	// 在 workerWg.Wait() 之后关闭，确保所有 worker 都已退出，不会再写入 errChan
	log.Debug("关闭错误通道")
	close(errChan)

	// 等待错误收集 goroutine 完成
	<-errCollected

	// 返回收集到的第一个错误（如果有）
	if processingErr != nil {
		log.Errorf("表处理过程中发生错误: %v", processingErr)
		return tablesTotal, processingErr
	}

	// 检查主上下文是否被取消
	if ctx.Err() != nil {
		log.Warnf("表处理过程被外部取消: %v", ctx.Err())
		return tablesTotal, ctx.Err()
	}

	log.Infof("成功调度 %d 个表进行处理", tablesTotal)
	return tablesTotal, nil
}

// 表信息结构
type tableInfo struct {
	Schema string
	Table  string
}

// TaskQueue 任务队列接口
//type TaskQueue interface {
//	Put(interface{}) error
//	Get() (interface{}, error)
//	Close() error
//	IsClosed() bool
//}

// createTaskQueue 创建任务队列
// *** 修改点: 确保返回值类型是 *queue.BoundedQueue 或实现了 queue.Queue 的其他类型 ***
func createTaskQueue() *queue.BoundedQueue { // 返回具体类型或接口 queue.Queue
	// 初始化任务队列，可以根据需求调整队列大小
	queueSize := 1000
	return queue.NewBoundedQueue(queueSize)
}

// waitForCompletionWithContext 支持上下文的表处理完成等待
func waitForCompletionWithContext(ctx context.Context, counter *atomic.Int32, expectedCount int) {
	checkInterval := 2 * time.Second

	log.Infof("等待 %d 张表完成处理...", expectedCount)

	ticker := time.NewTicker(checkInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			log.Warnf("等待表处理完成被取消: %v", ctx.Err())
			return
		case <-ticker.C:
			completedCount := counter.Load()
			if completedCount >= int32(expectedCount) {
				log.Infof("所有 %d 张表处理完成", completedCount)
				return
			}

			// 输出进度
			log.Infof("进度: %d/%d 表完成 (%.1f%%)",
				completedCount, expectedCount,
				float64(completedCount)/float64(expectedCount)*100)
		}
	}
}

// 统计配置中的表总数
func countConfiguredTables(tablesCheck map[string][]string) int {
	if tablesCheck == nil {
		return 0
	}

	count := 0
	for _, tables := range tablesCheck {
		count += len(tables)
	}
	return count
}

// 检查表在校验库中是否存在数据
func checkTableDataExists(checksumDB *sql.DB, schemaName, tableName, sourceDBID, targetDBID string) (sourceExists bool, targetExists bool, err error) {
	log.Debugf("检查表 %s.%s 在校验库中是否存在数据 (源DBID=%s, 目标DBID=%s)",
		schemaName, tableName, sourceDBID, targetDBID)

	if checksumDB == nil {
		log.Error("校验数据库连接为空")
		return false, false, fmt.Errorf("%w: 无效的校验数据库连接", ErrInvalidDBConnection)
	}

	// 检查源表数据
	log.Debugf("检查源表 %s.%s 数据是否存在 (DBID=%s)", schemaName, tableName, sourceDBID)
	sourceExists, err = compare.CheckTableExistsInChecksum(checksumDB, schemaName, tableName, sourceDBID)
	if err != nil {
		log.Debugf("检查源表数据存在性失败: %v", err)
	} else {
		log.Debugf("源表数据检查结果: %t", sourceExists)
	}

	// 检查目标表数据
	log.Debugf("检查目标表 %s.%s 数据是否存在 (DBID=%s)", schemaName, tableName, targetDBID)
	targetExists, err = compare.CheckTableExistsInChecksum(checksumDB, schemaName, tableName, targetDBID)
	if err != nil {
		log.Debugf("检查目标表数据存在性失败: %v", err)
	} else {
		log.Debugf("目标表数据检查结果: %t", targetExists)
	}

	// 确认检查结果
	if sourceExists && targetExists {
		log.Infof("表 %s.%s 源端和目标端的校验数据都存在", schemaName, tableName)
	} else {
		log.Warnf("表 %s.%s 校验数据不完整: 源端存在=%t, 目标端存在=%t",
			schemaName, tableName, sourceExists, targetExists)
	}

	return sourceExists, targetExists, err
}

// writeFixSQLToFile 写入修复SQL到文件
func writeFixSQLToFile(filePath string, sqls []string) error {
	if len(sqls) == 0 {
		log.Debugf("SQL语句为空，不需要写入文件: %s", filePath)
		return nil
	}

	// 获取目录和文件的绝对路径
	absPath, err := filepath.Abs(filePath)
	if err != nil {
		log.Warnf("获取文件绝对路径失败: %v, 将使用原始路径: %s", err, filePath)
		absPath = filePath
	} else {
		log.Debugf("生成SQL文件的绝对路径: %s", absPath)
	}

	// 确保目录存在
	dir := filepath.Dir(absPath)
	log.Debugf("确保SQL输出目录存在: %s", dir)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return fmt.Errorf("创建目录 %s 失败: %w", dir, err)
	}

	log.Debugf("开始创建文件: %s", absPath)
	file, err := os.Create(absPath)
	if err != nil {
		return fmt.Errorf("创建文件 %s 失败: %w", absPath, err)
	}
	defer func() {
		if closeErr := file.Close(); closeErr != nil {
			log.Debugf("关闭文件 %s 出错: %v", absPath, closeErr)
		}
	}()

	// 写入文件头
	targetID := config.GlobalConfig.DBTarget.ID()
	log.Debugf("写入SQL文件头 (目标DB ID: %s)", targetID)
	commentHeader := fmt.Sprintf("-- Auto-generated data fix SQL for target %s\n", targetID) +
		fmt.Sprintf("-- Generated at %s\n\n", time.Now().Format(time.RFC3339))

	if _, err := file.WriteString(commentHeader); err != nil {
		return fmt.Errorf("写入文件头到 %s 失败: %w", absPath, err)
	}

	// 写入SQL语句
	log.Debugf("开始写入 %d 条SQL语句到文件", len(sqls))
	for i, sqlStatement := range sqls {
		if _, err := file.WriteString(sqlStatement + "\n"); err != nil {
			return fmt.Errorf("写入第 %d 条 SQL 到 %s 失败: %w", i+1, absPath, err)
		}
	}

	// 确保数据已写入磁盘
	log.Debugf("执行文件同步，确保数据已写入磁盘")
	if err := file.Sync(); err != nil {
		log.Debugf("同步文件 %s 出错: %v", absPath, err)
	}

	fileInfo, err := file.Stat()
	if err != nil {
		log.Debugf("获取文件信息失败: %v", err)
	} else {
		log.Infof("成功写入修复SQL文件: %s (大小: %d 字节, %d 条语句)",
			absPath, fileInfo.Size(), len(sqls))
	}

	return nil
}

// cleanOldFixSQLFiles 清理指定表的旧修复SQL文件
func cleanOldFixSQLFiles(schemaName, tableName string) {
	if schemaName == "" || tableName == "" {
		log.Warn("清理修复SQL文件时模式名或表名为空")
		return
	}

	// 确保目录存在
	if _, err := os.Stat(outputDir); os.IsNotExist(err) {
		return // 目录不存在，不需要清理
	}

	// 查找匹配的文件模式
	pattern := filepath.Join(outputDir, fmt.Sprintf("datafix-%s.%s-*.sql", schemaName, tableName))
	matches, err := filepath.Glob(pattern)
	if err != nil {
		log.Warnf("查找旧修复文件失败 (%s.%s): %v", schemaName, tableName, err)
		return
	}

	if len(matches) == 0 {
		log.Debugf("没有找到与模式 %s 匹配的旧修复文件", pattern)
		return
	}

	// 删除匹配的文件
	deletedCount := 0
	for _, file := range matches {
		if err := os.Remove(file); err != nil {
			log.Warnf("删除旧修复文件失败 %s: %v", file, err)
		} else {
			deletedCount++
		}
	}

	if deletedCount > 0 {
		log.Debugf("已清理表 %s.%s 的 %d 个旧修复SQL文件", schemaName, tableName, deletedCount)
	}
}
